// Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::TensorBufferType;
use arrow::{array::{AsArray, UInt8Array}, datatypes::DataType};
use bytes::{Buf, Bytes};
use dora_core::config::DataId;
use dora_node_api::{ArrowData, Metadata};
use itertools::Itertools;
use nalgebra::{Const, DVector, Dyn, Matrix, RowVector3, ViewStorage, U1, U3};
use palette::{rgb::channels::Rgba, FromColor, Hsl, IntoColor, LinSrgb, Srgb};
use rerun::{components::{DepthMeter, DrawOrder, ViewCoordinates}, external::{arrow2::{self, buffer::Buffer}, re_log::error, re_types::ArrowString, re_types_core}, Archetype, AsComponents, Color, ComponentBatch, DepthImage, EntityPath, Image, LineStrips3D, MaybeOwnedComponentBatch, Pinhole, Points3D, Position3D, Radius, RecordingStream, RecordingStreamError, TensorBuffer, TensorData, TensorDimension, Text, TextLog, TextLogLevel, Vec2D, Vec3D};
use std::{collections::HashMap, mem, ops::Range};

pub struct DoraData {
    bytes: ArrowData,
    pub kind: DataKind,
}

const IMAGE_ENTITY: &str = "slam_viewer/images";
const TRAJECTORY_ENTITY: &str = "slam_viewer/trajectory";
const CAMERA_ENTITY: &str = "slam_viewer/camera_poses";
const RGB_ENTITY: &str = "slam_viewer/images/rgb";
const DEPTH_ENTITY: &str = "slam_viewer/images/depth";
const DEPTH_IMAGE_SCALING: f32 = 2.0;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DataKind {
    PointsXYZ,
    PointsXYZI,
    Image,
    Trajectory,
    TextLog,
    CameraPoses,
    RgbImage,
    DepthImage,
}

#[repr(C)]
#[derive(Clone, Copy)]
pub union Points3DOrIntensity {
    position: Position3D,
    intensity: f32,
}

pub struct Points3DOwned {
    pub positions: Vec<Position3D>,
    pub colors: Vec<Color>,
    pub radii: Vec<Radius>,
}

pub struct Points3DRef<'a> {
    positions: &'a [Position3D],
    colors: Vec<Color>,
    radii: Vec<Radius>,
}

pub struct TextLogRef<'a> {
    bytes: &'a [u8],
    text_range: Range<usize>,
    level: Option<TextLogLevel>,
}

#[derive(Debug, thiserror::Error)]
pub enum DataKindConvertError<'a> {
    #[error("unsupported data type: {0}")]
    UnsupportedDataType(&'a DataType),
    #[error("unknown input id: {0}")]
    UnknownInputId(&'a DataId),
}

type PointsViewXYZ<'a, T> = ViewStorage<'a, T, Dyn, U3, U3, U1>;
type PointsMatXYZ<'a, T> = Matrix<T, Dyn, U3, PointsViewXYZ<'a, T>>;

macro_rules! tensor_buffer_from_bytes {
    ($bytes:expr, $buffer_type:expr, $($variant:ident => $ty:ty),+) => {
        match $buffer_type {
            $(
                Ok(TensorBufferType::$variant) => {
                    if $bytes.len() % size_of::<$ty>() != 0 {
                        panic!("Buffer length is not a multiple of element size for type {}", stringify!($ty));
                    }
                    let slice = unsafe {
                        std::slice::from_raw_parts(
                            $bytes.as_ptr() as *const $ty,
                            $bytes.len() / size_of::<$ty>()
                        ).to_vec()
                    };
                    TensorBuffer::$variant(re_types_core::ArrowBuffer::<$ty>::from(slice))
                },
            )+
            _ => unimplemented!("Unsupported buffer type {:?} ", $buffer_type),
        }
    };
}

fn buf_to_tensor_data(mut bytes: Bytes) -> TensorData {
    let _nbytes = bytes.get_u64() as usize;
    let mut shape = Vec::new();
    let shape_num = bytes.get_u64() as usize;
    for _ in 0..shape_num {
        let dimension_size = bytes.get_u64();
        shape.push(TensorDimension {
            size: dimension_size,
            name: None,
        });
    }
    let tensor_buffer_type = bytes.get_u8();
    let buffer_type = TensorBufferType::try_from(tensor_buffer_type);
    let tensor_buffer = tensor_buffer_from_bytes!(
        bytes,
        buffer_type,
        U8 => u8,
        U16 => u16,
        U32 => u32,
        U64 => u64,
        I8 => i8,
        I16 => i16,
        I32 => i32,
        I64 => i64,
        F16 => arrow2::types::f16,
        F32 => f32,
        F64 => f64,
        Jpeg => u8,
        Nv12 => u8,
        Yuy2 => u8
    );
    TensorData {
        shape,
        buffer: tensor_buffer,
    }
}

fn buf_to_tensor_data_with_shape_name(mut bytes: Bytes) -> TensorData {
    let _nbytes = bytes.get_u64() as usize;
    let mut shape = Vec::new();
    let shape_num = bytes.get_u64() as usize;
    for _ in 0..shape_num {
        let dimension_size = bytes.get_u64();
        let name_length = bytes.get_u64() as usize;
        if name_length > 0 && bytes.remaining() >= name_length {
            let name_bytes = bytes.split_to(name_length).to_vec();
            let dimension_name = Some(ArrowString(Buffer::from(name_bytes)));
            shape.push(TensorDimension {
                size: dimension_size,
                name: dimension_name,
            });
        } else {
            shape.push(TensorDimension {
                size: dimension_size,
                name: None,
            });
        }
    }
    let tensor_buffer_type = bytes.get_u8();
    let buffer_type = TensorBufferType::try_from(tensor_buffer_type);
    let tensor_buffer = tensor_buffer_from_bytes!(
        bytes,
        buffer_type,
        U8 => u8,
        U16 => u16,
        U32 => u32,
        U64 => u64,
        I8 => i8,
        I16 => i16,
        I32 => i32,
        I64 => i64,
        F16 => arrow2::types::f16,
        F32 => f32,
        F64 => f64,
        Jpeg => u8,
        Nv12 => u8,
        Yuy2 => u8
    );
    TensorData {
        shape,
        buffer: tensor_buffer,
    }
}

impl DoraData {
    pub fn new(bytes: ArrowData, kind: DataKind) -> Self {
        Self { bytes, kind }
    }

    pub fn log(
        &self,
        recording_stream: &RecordingStream,
        ent_path: impl Into<EntityPath>,
    ) -> Result<(), RecordingStreamError> {
        match self.kind {
            DataKind::PointsXYZ => recording_stream.log(ent_path, &self.as_points_xyz()),
            DataKind::PointsXYZI => recording_stream.log(ent_path, &self.as_points_xyzi()),
            DataKind::TextLog => recording_stream.log(ent_path, &self.as_text_log()),
            DataKind::Image => recording_stream.log(IMAGE_ENTITY, &self.to_image()),
            DataKind::Trajectory => recording_stream.log(TRAJECTORY_ENTITY, &self.to_trajectory()),
            DataKind::CameraPoses => {
                recording_stream.log(CAMERA_ENTITY, &self.to_camera_poses())?;
                recording_stream.log(CAMERA_ENTITY, &self.to_camera_poses_image())?;
                Ok(())
            },
            DataKind::RgbImage => recording_stream.log(RGB_ENTITY, &self.to_camera_poses_image()),
            DataKind::DepthImage => {
                recording_stream.log(IMAGE_ENTITY, &self.to_depth_camera_poses())?;
                recording_stream.log(DEPTH_ENTITY, &self.to_depth_image())?;
                Ok(())
            },
        }
    }

    pub fn as_points_xyz(&self) -> Points3DRef<'_> {
        debug_assert_eq!(self.kind, DataKind::PointsXYZ);
        let array: &UInt8Array = self.bytes.as_primitive();
        let positions = unsafe {
            std::slice::from_raw_parts(
                array.values().as_ptr() as *const Position3D,
                self.bytes.len() / size_of::<Position3D>(),
            )
        };
        let (colors, radii) = gen_points3d_components(positions);
        Points3DRef {
            positions,
            colors,
            radii,
        }
    }

    pub fn as_points_xyzi(&self) -> Points3DOwned {
        debug_assert_eq!(self.kind, DataKind::PointsXYZI);
        let array: &UInt8Array = self.bytes.as_primitive();
        let size_of_element = mem::size_of::<Position3D>() + mem::size_of::<f32>();
        let num_elements = array.values().len() / size_of_element;
        let mut positions = Vec::with_capacity(num_elements);
        let mut intensities = Vec::with_capacity(num_elements);
        let points_with_intensity = unsafe {
            std::slice::from_raw_parts(
                array.values().as_ptr() as *const u8,
                self.bytes.len() / size_of::<u8>(),
            )
        };
        for i in 0..num_elements {
            let point_start = i * size_of_element;
            let position = unsafe {
                *(points_with_intensity[point_start..point_start + mem::size_of::<Position3D>()]
                    .as_ptr() as *const Position3D)
            };
            let intensity = unsafe {
                *(points_with_intensity
                    [point_start + mem::size_of::<Position3D>()..point_start + size_of_element]
                    .as_ptr() as *const f32)
            };
            positions.push(position);
            intensities.push(intensity);
        }
        let (colors, radii) = gen_points3d_components_with_intensity(&positions, intensities);
        Points3DOwned {
            positions,
            colors,
            radii,
        }
    }

    pub fn as_text_log(&self) -> TextLogRef<'_> {
        debug_assert_eq!(self.kind, DataKind::TextLog);
        let array: &UInt8Array = self.bytes.as_primitive();
        let bytes = array.values().inner().as_slice();
        let str = String::from_utf8_lossy(bytes);
        let mut text = str.as_ref();
        let mut level = None;
        if let Some(s) = text.strip_prefix(TextLogLevel::CRITICAL) {
            text = s;
            level = Some(TextLogLevel::CRITICAL);
        } else if let Some(s) = text.strip_prefix(TextLogLevel::ERROR) {
            text = s;
            level = Some(TextLogLevel::ERROR);
        } else if let Some(s) = text.strip_prefix(TextLogLevel::WARN) {
            text = s;
            level = Some(TextLogLevel::WARN);
        } else if let Some(s) = text.strip_prefix(TextLogLevel::INFO) {
            text = s;
            level = Some(TextLogLevel::INFO);
        } else if let Some(s) = text.strip_prefix(TextLogLevel::DEBUG) {
            text = s;
            level = Some(TextLogLevel::DEBUG);
        } else if let Some(s) = text.strip_prefix(TextLogLevel::TRACE) {
            text = s;
            level = Some(TextLogLevel::TRACE);
        }
        text = text.trim();
        let offset = text.as_ptr() as usize - str.as_ptr() as usize;

        TextLogRef {
            bytes,
            text_range: offset..offset + text.len(),
            level: level.map(|l| TextLogLevel::from(l)),
        }
    }

    pub fn to_image(&self) -> Image {
        debug_assert_eq!(self.kind, DataKind::Image);
        let array: &UInt8Array = self.bytes.as_primitive();
        let bytes = array.values().inner().as_slice();
        Image {
            data: rerun::components::TensorData(buf_to_tensor_data(Bytes::copy_from_slice(bytes))),
            draw_order: Some(DrawOrder(0.0)),
        }
    }

    pub fn to_trajectory(&self) -> LineStrips3D {
        debug_assert_eq!(self.kind, DataKind::Trajectory);
        let array: &UInt8Array = self.bytes.as_primitive();
        let strips = unsafe {
            std::slice::from_raw_parts(
                array.values().as_ptr() as *const Vec3D,
                self.bytes.len() / size_of::<Vec3D>(),
            )
        };
        LineStrips3D::new(vec![strips]).with_radii([Radius(0.03)])
    }

    pub fn to_camera_poses_image(&self) -> Image {
        let array: &UInt8Array = self.bytes.as_primitive();
        let bytes = array.values().inner().as_slice();
        Image {
            data: rerun::components::TensorData(buf_to_tensor_data_with_shape_name(
                Bytes::copy_from_slice(bytes),
            )),
            draw_order: Some(DrawOrder(0.0)),
        }
    }

    pub fn to_camera_poses(&self) -> Pinhole {
        debug_assert_eq!(self.kind, DataKind::CameraPoses);
        let array: &UInt8Array = self.bytes.as_primitive();
        let bytes = array.values().inner().as_slice();
        let image = Image {
            data: rerun::components::TensorData(buf_to_tensor_data_with_shape_name(
                Bytes::copy_from_slice(bytes),
            )),
            draw_order: Some(DrawOrder(0.0)),
        };
        let mut dimensions = HashMap::new();
        let shape = image.data.shape();
        for dim in shape.iter() {
            if let Some(name) = &dim.name {
                match name.as_str() {
                    "height" => {
                        dimensions.insert("height", dim.size);
                    }
                    "width" => {
                        dimensions.insert("width", dim.size);
                    }
                    _ => {}
                }
            }
        }
        if let (Some(&height), Some(&width)) = (dimensions.get("height"), dimensions.get("width")) {
            let focal_length = Vec2D::new(width as f32 * 0.7, width as f32 * 0.7);
            let resolution = Vec2D::new(width as f32, height as f32);
            let pinhole = Pinhole::from_focal_length_and_resolution(focal_length, resolution);
            return pinhole;
        } else {
            eprintln!("Error: Could not find both 'height' and 'width' in tensor dimensions.");
            let pinhole = Pinhole::from_focal_length_and_resolution([0.0, 0.0], (0.0, 0.0));
            return pinhole;
        }
    }

    pub fn to_depth_camera_poses(&self) -> Pinhole {
        let array: &UInt8Array = self.bytes.as_primitive();
        let bytes = array.values().inner().as_slice();
        let image = Image {
            data: rerun::components::TensorData(buf_to_tensor_data_with_shape_name(
                Bytes::copy_from_slice(bytes),
            )),
            draw_order: Some(DrawOrder(0.0)),
        };
        let mut dimensions = HashMap::new();
        let shape = image.data.shape();
        for dim in shape.iter() {
            if let Some(name) = &dim.name {
                match name.as_str() {
                    "height" => {
                        dimensions.insert("height", dim.size);
                    }
                    "width" => {
                        dimensions.insert("width", dim.size);
                    }
                    _ => {}
                }
            }
        }
        if let (Some(&height), Some(&width)) = (dimensions.get("height"), dimensions.get("width")) {
            let focal_length = Vec2D::new(width as f32 * 0.7, width as f32 * 0.7);
            let resolution = Vec2D::new(width as f32, height as f32);
            let pinhole = Pinhole::from_focal_length_and_resolution(focal_length, resolution)
                .with_camera_xyz(ViewCoordinates::RDF);
            return pinhole;
        } else {
            eprintln!("Error: Could not find both 'height' and 'width' in tensor dimensions.");
            let pinhole = Pinhole::from_focal_length_and_resolution([0.0, 0.0], (0.0, 0.0));
            return pinhole;
        }
    }

    pub fn to_depth_image(&self) -> DepthImage {
        let array: &UInt8Array = self.bytes.as_primitive();
        let bytes = array.values().inner().as_slice();
        let image = Image {
            data: rerun::components::TensorData(buf_to_tensor_data_with_shape_name(
                Bytes::copy_from_slice(bytes),
            )),
            draw_order: Some(DrawOrder(0.0)),
        };
        DepthImage {
            data: image.data,
            meter: Some(DepthMeter(DEPTH_IMAGE_SCALING)),
            draw_order: image.draw_order,
        }
    }
}

impl DataKind {
    pub fn try_from_input_event<'a>(
        id: &'a DataId,
        metadata: &'a Metadata,
    ) -> Result<Self, DataKindConvertError<'a>> {
        let data_type = &metadata.type_info.data_type;
        if !matches!(data_type, DataType::UInt8) {
            return Err(DataKindConvertError::UnsupportedDataType(data_type));
        }
        match id.as_str() {
            "points_xyz" => Ok(Self::PointsXYZ),
            "points_xyzi" => Ok(Self::PointsXYZI),
            "text_log" => Ok(Self::TextLog),
            "image" => Ok(Self::Image),
            "trajectory" => Ok(Self::Trajectory),
            "camera_poses" => Ok(Self::CameraPoses),
            "rgb_image" => Ok(Self::RgbImage),
            "depth_image" => Ok(Self::DepthImage),
            _ => Err(DataKindConvertError::UnknownInputId(id)),
        }
    }

    pub fn as_output_id(&self) -> &'static str {
        match self {
            DataKind::PointsXYZ => "points_xyz",
            DataKind::PointsXYZI => "points_xyzi",
            DataKind::TextLog => "text_log",
            DataKind::Image => "image",
            DataKind::Trajectory => "trajectory",
            DataKind::CameraPoses => "camera_poses",
            DataKind::RgbImage => "rgb_image",
            DataKind::DepthImage => "depth_image",
        }
    }
}

impl<'a> Points3DRef<'a> {
    const DEFAULT_RADIUS: Radius = Radius(0.001);
}

impl<'a> AsComponents for Points3DRef<'a> {
    fn as_component_batches(&self) -> Vec<MaybeOwnedComponentBatch<'_>> {
        [
            Points3D::indicator(),
            MaybeOwnedComponentBatch::Ref(&self.positions),
            MaybeOwnedComponentBatch::Ref(&self.colors),
            MaybeOwnedComponentBatch::Ref(&self.radii),
        ]
        .into_iter()
        .collect()
    }
}

impl<'a> AsComponents for Points3DOwned {
    fn as_component_batches(&self) -> Vec<MaybeOwnedComponentBatch<'_>> {
        [
            Points3D::indicator(),
            MaybeOwnedComponentBatch::Ref(&self.positions),
            MaybeOwnedComponentBatch::Ref(&self.colors),
            MaybeOwnedComponentBatch::Ref(&self.radii),
        ]
        .into_iter()
        .collect()
    }
}

impl<'a> AsComponents for TextLogRef<'a> {
    fn as_component_batches(&self) -> Vec<MaybeOwnedComponentBatch<'_>> {
        let s = String::from_utf8_lossy(&self.bytes[self.text_range.start..self.text_range.end]);
        let text = Text::from(s.as_ref());
        [
            Some(TextLog::indicator()),
            Some(MaybeOwnedComponentBatch::Owned(Box::new(text))),
            self.level
                .as_ref()
                .map(|l| (l as &dyn ComponentBatch).into()),
        ]
        .into_iter()
        .flatten()
        .collect()
    }
}

pub fn gen_points3d_components(positions: &[Position3D]) -> (Vec<Color>, Vec<Radius>) {
    let points = Matrix::from_data(unsafe {
        ViewStorage::from_raw_parts(
            positions.as_ptr() as *const f32,
            (Dyn(positions.len()), Const::<3>),
            (Const::<3>, Const::<1>),
        )
    });
    let center = RowVector3::new(0.0, 0.0, 0.0);
    let colors = gen_colors(points, center);
    let radii = vec![Points3DRef::DEFAULT_RADIUS; positions.len()];
    (colors, radii)
}

fn hsl_to_vec(hsl: Hsl) -> RowVector3<f32> {
    RowVector3::new(hsl.hue.into_inner(), hsl.saturation, hsl.lightness)
}

fn gen_colors<'a>(points: PointsMatXYZ<'a, f32>, center: RowVector3<f32>) -> Vec<Color> {
    let color_from = palette::named::PALEVIOLETRED;
    let color_to = palette::named::PALETURQUOISE;
    let color_from_hsl = Hsl::from_color(color_from.into_format::<f32>());
    let color_to_hsl = Hsl::from_color(color_to.into_format::<f32>());
    let color_diff = color_to_hsl - color_from_hsl;
    let color_from_vec = hsl_to_vec(color_from_hsl);
    let color_diff_vec = hsl_to_vec(color_diff);
    let distances = DVector::<f32>::from_iterator(
        points.nrows(),
        points.row_iter().map(|r| (r - center).norm()),
    );
    let max_distance = distances.max();
    let ratios = distances / max_distance;
    let color_vars = ratios * color_diff_vec;
    color_vars
        .row_iter()
        .map(|c| c + color_from_vec)
        .map(|c| {
            let hsl = Hsl::new_srgb(c.x, c.y, c.z);
            let srgb: Srgb = hsl.into_color();
            Color::from_u32(srgb.into_format::<u8>().into_u32::<Rgba>())
        })
        .collect()
}

pub fn gen_points3d_components_with_intensity(
    positions: &[Position3D],
    intensity: Vec<f32>,
) -> (Vec<Color>, Vec<Radius>) {
    let colors = gen_colors_with_intensity(intensity);
    let radii = vec![Points3DRef::DEFAULT_RADIUS; positions.len()];
    (colors, radii)
}

fn gen_colors_with_intensity<'a>(intensity: Vec<f32>) -> Vec<Color> {
    let color_from = palette::named::PALEVIOLETRED;
    let color_to = palette::named::PALETURQUOISE;
    let color_from_hsl = Hsl::from_color(color_from.into_format::<f32>());
    let color_to_hsl = Hsl::from_color(color_to.into_format::<f32>());
    let (min_intensity, max_intensity) = match intensity.iter().minmax() {
        itertools::MinMaxResult::NoElements => (f32::INFINITY, f32::NEG_INFINITY),
        itertools::MinMaxResult::OneElement(&e) => (e, e),
        itertools::MinMaxResult::MinMax(&min, &max) => (min, max),
    };
    let intensity_range = max_intensity - min_intensity;
    let intensity_range = if intensity_range == 0.0 {
        1.0
    } else {
        intensity_range
    };
    intensity
        .iter()
        .map(|&i| {
            let ratio = (i as f32 - min_intensity) / intensity_range;
            let hsl = Hsl::new(
                color_from_hsl.hue.into_positive_degrees()
                    + ratio
                        * (color_to_hsl.hue.into_positive_degrees()
                            - color_from_hsl.hue.into_positive_degrees()),
                color_from_hsl.saturation
                    + ratio * (color_to_hsl.saturation - color_from_hsl.saturation),
                color_from_hsl.lightness
                    + ratio * (color_to_hsl.lightness - color_from_hsl.lightness),
            );
            let lin_srgb: LinSrgb = hsl.into_color();
            let srgb: Srgb<f32> = lin_srgb.into();
            Color::from_u32(srgb.into_format::<u8>().into_u32::<Rgba>())
        })
        .collect()
}